package cn.edu.sdjzu.knoverse.service;

import cn.edu.sdjzu.knoverse.utils.rag.CustomQueryRouter;
import cn.edu.sdjzu.knoverse.utils.rag.EnhancedVectorRetriever;
import cn.edu.sdjzu.knoverse.utils.rag.RagAssistant;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.neo4j.Neo4jContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.rag.query.router.QueryRouter;
import dev.langchain4j.rag.query.transformer.CompressingQueryTransformer;
import dev.langchain4j.rag.query.transformer.QueryTransformer;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.Result;
import dev.langchain4j.store.embedding.EmbeddingStore;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;

@Service
public class RagService {
    private static final Logger log = LoggerFactory.getLogger(RagService.class);

    @Resource
    private ChatLanguageModel chatLanguageModel;

    @Resource
    private StreamingChatLanguageModel streamingChatLanguageModel;

    @Resource
    private EmbeddingModel embeddingModel;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    @Resource(name = "vectorContentRetriever")
    private ContentRetriever vectorContentRetriever;

    @Resource
    private Neo4jContentRetriever neo4jContentRetriever;

    // 新增增强型检索器
    private ContentRetriever enhancedVectorRetriever;

    private ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

    // 缓存查询增强组件
    private QueryTransformer queryTransformer;
    private QueryRouter queryRouter;
    private RetrievalAugmentor retrievalAugmentor;

    /**
     * RAG增强的流式聊天，先检索相关内容，再使用流式模型回复
     *
     * @param message 用户消息
     * @param retrievalType 检索方式：VECTOR（向量数据库）、GRAPH（图数据库）、ENHANCED（图增强向量检索）、PARALLEL（并行检索）或AUTO（自动决定）
     * @param handler 流式响应处理器
     */
    public void streamRagChat(String message, String retrievalType, CustomStreamingChatResponseHandler handler) {
        log.info("处理流式RAG聊天请求: {}, 检索方式: {}", message, retrievalType);

        try {
            // 初始化检索组件
            prepareRetrievalComponents();

            // 创建查询，不使用查询压缩转换器
            Query query = Query.from(message);

            // 根据检索类型选择检索器
            List<Content> retrievedContents = new ArrayList<>();

            // 如果是AUTO，则由查询路由器决定使用哪些检索器
            if ("AUTO".equalsIgnoreCase(retrievalType)) {
                Collection<ContentRetriever> contentRetrievers = queryRouter.route(query);
                
                // 如果不需要检索，直接使用流式模型
                if (contentRetrievers.isEmpty()) {
                    log.info("跳过检索，直接使用流式模型回复");
                    // 发送跳过检索的状态
                    handler.onStatus("检索状态", "跳过检索");
                    streamingChatLanguageModel.chat(message, handler);
                    return;
                }
                
                // 执行检索
                for (ContentRetriever retriever : contentRetrievers) {
                    // 根据检索器类型发送不同的状态消息
                    sendRetrieverStatusMessage(handler, retriever);
                    
                    List<Content> contents = retriever.retrieve(query);
                    if (contents != null && !contents.isEmpty()) {
                        retrievedContents.addAll(contents);
                    }
                }
            } else {
                // 根据指定的检索类型执行检索
                List<ContentRetriever> selectedRetrievers = getRetrieversForType(retrievalType);
                for (ContentRetriever retriever : selectedRetrievers) {
                    // 发送检索状态
                    sendRetrieverStatusMessage(handler, retriever);
                    
                    List<Content> contents = retriever.retrieve(query);
                    if (contents != null && !contents.isEmpty()) {
                        retrievedContents.addAll(contents);
                    }
                }
            }

            // 发送检索完成的状态
            handler.onStatus("检索状态", "检索完成");

            // 构建包含检索内容的提示
            String retrievedInfo = retrievedContents.stream()
                    .map(content -> content.textSegment().text())
                    .collect(Collectors.joining("\n\n"));

            // 构建RAG提示模板
            String promptText = "您是一个知识助手。基于以下从图数据库或向量数据库中检索到的信息回答用户的问题。"
                    + "根据用户问题检索到的信息:\n``` " + retrievedInfo + " ```\n\n"
                    + "如果查询出的信息为空或者无法从信息中找到答案，请基于您的知识回答，但要明确表示没有从知识库中检索到相关信息。\n\n"
                    + "用户问题: ``` " + message+" ```";

            log.info("构建的RAG提示: {}", promptText);
            // 发送模型思考中的状态
            handler.onStatus("检索状态", "生成回答中");
            
            // 使用流式模型回复
            streamingChatLanguageModel.chat(promptText, handler);

        } catch (Exception e) {
            log.error("流式RAG聊天出错", e);
            handler.onError(e);
        }
    }

    /**
     * 根据检索类型获取相应的检索器
     * 
     * @param retrievalType 检索类型
     * @return 检索器列表
     */
    private List<ContentRetriever> getRetrieversForType(String retrievalType) {
        List<ContentRetriever> retrievers = new ArrayList<>();
        
        switch (retrievalType.toUpperCase()) {
            case "VECTOR":
                retrievers.add(vectorContentRetriever);
                break;
            case "GRAPH":
                retrievers.add(neo4jContentRetriever);
                break;
            case "ENHANCED":
                retrievers.add(enhancedVectorRetriever);
                break;
            case "PARALLEL":
                retrievers.add(vectorContentRetriever);
                retrievers.add(neo4jContentRetriever);
                break;
            default:
                // 默认使用AUTO方式，这里不应该执行到
                log.warn("未知的检索类型: {}, 将由查询路由器决定", retrievalType);
                break;
        }
        
        return retrievers;
    }
    
    /**
     * 根据检索器类型发送相应的状态消息
     * 
     * @param handler 响应处理器
     * @param retriever 检索器
     */
    private void sendRetrieverStatusMessage(CustomStreamingChatResponseHandler handler, ContentRetriever retriever) {
        if (retriever == neo4jContentRetriever) {
            handler.onStatus("检索状态", "图谱查询中");
        } else if (retriever == vectorContentRetriever) {
            handler.onStatus("检索状态", "向量召回中");
        } else if (retriever == enhancedVectorRetriever) {
            handler.onStatus("检索状态", "图增强向量检索中");
        } else {
            handler.onStatus("检索状态", "正在检索中");
        }
    }
    /**
     * 准备检索相关组件
     */
    private void prepareRetrievalComponents() {
        // 创建查询转换器（如果尚未创建）- 用于查询压缩
        if (queryTransformer == null) {
            queryTransformer = new CompressingQueryTransformer(chatLanguageModel);
            log.info("创建了查询压缩转换器");
        }

        // 创建增强型检索器（如果尚未创建）
        if (enhancedVectorRetriever == null) {
            enhancedVectorRetriever = new EnhancedVectorRetriever(
                    neo4jContentRetriever,
                    vectorContentRetriever,
                    chatLanguageModel
            );
            log.info("创建了增强型检索器");
        }

        // 创建查询路由器（如果尚未创建）- 智能选择检索器和跳过检索
        if (queryRouter == null) {
            // 使用自定义查询路由器，支持跳过检索和智能选择检索器
            queryRouter = new CustomQueryRouter(
                    chatLanguageModel,
                    vectorContentRetriever,
                    neo4jContentRetriever,
                    enhancedVectorRetriever  // 添加增强型检索器
            );
            log.info("创建了自定义查询路由器");
        }
    }

    /**
     * 使用RAG进行对话，并返回检索来源
     *
     * @param message 用户消息
     * @return 包含回复内容和来源的Result对象
     */
    public Result<String> chatWithSources(String message) {
        // 创建查询转换器（如果尚未创建）- 用于查询压缩
        if (queryTransformer == null) {
            queryTransformer = new CompressingQueryTransformer(chatLanguageModel);
            log.info("创建了查询压缩转换器");
        }

        // 创建增强型检索器（如果尚未创建）
        if (enhancedVectorRetriever == null) {
            enhancedVectorRetriever = new EnhancedVectorRetriever(
                    neo4jContentRetriever,
                    vectorContentRetriever,
                    chatLanguageModel
            );
            log.info("创建了增强型检索器");
        }

        // 创建查询路由器（如果尚未创建）- 智能选择检索器和跳过检索
        if (queryRouter == null) {
            // 使用自定义查询路由器，支持跳过检索和智能选择检索器
            queryRouter = new CustomQueryRouter(
                    chatLanguageModel,
                    vectorContentRetriever,
                    neo4jContentRetriever,
                    enhancedVectorRetriever  // 添加增强型检索器
            );
            log.info("创建了自定义查询路由器");
        }

        // 创建检索增强器（如果尚未创建）
        if (retrievalAugmentor == null) {
            retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                    .queryTransformer(queryTransformer) // 添加查询压缩
                    .queryRouter(queryRouter) // 添加查询路由和跳过检索
                    .build();
            log.info("创建了检索增强器");
        }

        // 创建AI服务
        RagAssistant assistant = AiServices.builder(RagAssistant.class)
                .chatLanguageModel(chatLanguageModel)
                .retrievalAugmentor(retrievalAugmentor)
                .chatMemory(chatMemory)
                .build();

        // 调用AI服务进行对话
        log.info("处理用户查询: {}", message);
        return assistant.chat(message);
    }

    /**
     * 处理文档并存储到向量数据库
     *
     * @param document 文档内容
     * @param fileName 文件名
     */
    public void processDocument(String document, String fileName) {
        // 解析文档
        DocumentParser documentParser = new TextDocumentParser();

        // 创建Metadata对象，包含文件名和上传时间
        String uploadTime = LocalDateTime.now().format(DateTimeFormatter.ISO_LOCAL_DATE_TIME);
        String documentId = UUID.randomUUID().toString();

        // 使用Map创建Metadata
        Map<String, String> metadataMap = new HashMap<>();
        metadataMap.put("source", fileName);
        metadataMap.put("upload_time", uploadTime);
        metadataMap.put("document_id", documentId);

        Metadata docMetadata = Metadata.from(metadataMap);

        // 创建Document对象
        Document doc = Document.from(document, docMetadata);

        // 分割文档
        DocumentSplitter splitter = DocumentSplitters.recursive(300, 0);
        List<TextSegment> segments = splitter.split(doc);

        // 生成嵌入向量
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

        // 存储到向量数据库
        embeddingStore.addAll(embeddings, segments);
        log.info("成功处理并存储文档: {}, 共{}个文本段落", fileName, segments.size());
    }
}

